Consider the benchmark

using JuMP

f(x, u) = [sin(x[1]) - x[1] * u, cos(x[2]) + x[1] * u]
function RK4(f, X, u)
    k1 = f(X     , u)
    k2 = f(X+k1/2, u)
    k3 = f(X+k2/2, u)
    k4 = f(X+k3  , u)
    X + (k1 + 2k2 + 2k3 + k4) / 6
end

import Ipopt
function bench(n)
    model = direct_model(Ipopt.Optimizer())

    @variable(model, q[1:2])
    @variable(model, u)

    x = q
    for _ = 1:n
        x = RK4(f, x, u)
    end

    @constraint(model, x .== 0);

    @objective(model, Min, u^2)

    @time optimize!(model)
end
@time bench(4)

When combined with jump-dev/JuMP.jl#4032, I now get:

  0.002103 seconds (3.79 k allocations: 226.328 KiB)
  1.121932 seconds (9.58 k allocations: 513.109 KiB)

All the time is actually being spent by check_belongs_to_model in https://github.com/jump-dev/JuMP.jl/blob/b19c3e71e74e87bbf51d84cb1cfb94b0d8e42700/src/constraints.jl#L1037, if I comment out this line, I get:

  0.002041 seconds (3.79 k allocations: 226.328 KiB)
  0.002394 seconds (9.55 k allocations: 512.219 KiB)

This could be fixed by checking the model at the same time as we do moi_function.

Another caveat is that we need to use direct_model, otherwise, map_indices in

I fixed the performance issue of check_belongs_to_model and the issue of map_indices creating duplicates, it works now even without direct model

using JuMP

f(x, u) = [sin(x[1]) - x[1] * u, cos(x[2]) + x[1] * u]
function RK4(f, X, u)
    k1 = f(X     , u)
    k2 = f(X+k1/2, u)
    k3 = f(X+k2/2, u)
    k4 = f(X+k3  , u)
    X + (k1 + 2k2 + 2k3 + k4) / 6
end

import Ipopt
function bench(n)
    model = Model(Ipopt.Optimizer)

    @variable(model, q[1:2])
    @variable(model, u)

    x = q
    for _ = 1:n
        x = RK4(f, x, u)
    end

    @constraint(model, x .== 0);

    @objective(model, Min, u^2)

    @time optimize!(model)
end
bench(1)
bench(2)
bench(3)
@time bench(4)

I get

  0.001918 seconds (6.61 k allocations: 426.477 KiB)
  0.002599 seconds (11.97 k allocations: 614.852 KiB)

There's a cost to adding subexpressions. I need to think very carefully if we should do this. It's not an obvious win. I've been procrastinating on this, not because it is technically difficult, but because I'm not sure if it is something that MOI should even do. Users can manually extract their subexpressions if they desire.

I think at minimum, we're going to need a much larger set of benchmarks.

I agree, it's not a clear choice. What convinced me to write this PR is the following. When users share sub-expressions by reference, we either:

1. Exponentially blow up the problem size (current)
2. Handle it proportionally (this PR)

The exponential blowup can make problems intractable. We have workarounds but most users won't think about them, they will just have their computer freeze. On the other hand, having more sub-expressions than necessary can at worse be a bit slower. And we're not inventing sub-expressions, the user already had them in their code, either unintentionally or for saving memory.

Post-hoc detection (like what you did in https://github.com/lanl-ansi/MathOptSymbolicAD.jl) is complementary as it catches more sub-expressions but it doesn't prevent the blowup during construction.

I am happy if this works only in direct-mode. I need to think about this more. If we can build the tape fast, but it's slow to evaluate, then I'd consider writing a new AD backend that is structured as a complete DAG, rather than the current function-based approach.

Do we have a benchmark showing that the subexpressions are slow ? I think you shared one at some point but I don't remember if it showed that it was slower or just comparable so not worth it. We can definitely split this PR in smaller ones.

This part should be enough for direct mode, the other parts were moved to #2802 and #2803